#include "structMeshIdxTranslation.hpp"

namespace hms {


template<typename MeshType>
void setLocalPatch(
	BoundaryPatch& patch,
	StructMeshPart<MeshType>& part,
	const MeshType& mesh,
	const Ref<const Array2i>& globalEdgeLimits
){
	/* translate global edge indices to local ones */
	BoundaryLocation loc { whereIs(globalEdgeLimits(0), part.meshSizeGlobal()) };

	bool isHorizontal {
		loc == BoundaryLocation::bottom || loc == BoundaryLocation::top
	};
	
	auto getLocalEdges = [&](
		Index stride, auto* localToGlobal, auto* globalToLocal
	) -> ArrayXi {
		
		auto toGlobal = [&]( auto* getEdge ){
			return localToGlobal( getEdge( loc, part.meshSize() ),
				part.meshSizeGlobal(), part.meshSize(), part.offset()
			);
		};
		std::optional<Array2i> limits { overlap(
			globalEdgeLimits,
			Array2i{ toGlobal(firstEdge), toGlobal(lastEdge) }
		)};
		if ( limits )
			return ArrayXi::LinSpaced(
				( (*limits)[1]-(*limits)[0] ) / stride + 1,
				globalToLocal(
					(*limits)[0],
					part.meshSizeGlobal(), part.meshSize(), part.offset()
				),
				globalToLocal(
					(*limits)[1],
					part.meshSizeGlobal(), part.meshSize(), part.offset()
				)
			);
		else
			return {};
	};
	
	if ( part.atBoundary(loc) ){
		if (isHorizontal)
			patch.setEdgesAndCells( mesh, getLocalEdges(
				1, /* stride */
				hEdgeLocalToGlobal,
				hEdgeGlobalToLocal
			) );
		else
			patch.setEdgesAndCells( mesh, getLocalEdges(
				part.meshSizeGlobal().x() + 1, /* stride */
				vEdgeLocalToGlobal,
				vEdgeGlobalToLocal
			) );
	}
	else {
		/* set patch as empty */
		patch.clearIndices();
	}
}



template<typename MeshType>
Array2Xi allgatherBoundaryPatchLimits( const StructMeshPart<MeshType>& part ){
	/* just a shorthand */
	const std::vector<BoundaryPatch>& patches {
		part.mesh().boundaryPatches()
	};
	auto nPatches { patches.size() };
	/* assumption is that all mesh partitions already have all boundary patches.
	 * To validate that, we overwrite nPatches from master and check. */
	MPI_Bcast_H( nPatches, part.comm(), part.comm().master() );
	if ( nPatches != patches.size() ){
		std::cerr <<
		"Number of boundary patches must be equal on all mesh partitions!\n";
		MPI_Abort(part.comm(), MPI_ERR_OTHER);
	}
	/* we also need to validate that the patches have the same names on all
	 * partitions, otherwise accessing them through names will be an issue */
	for ( const auto& patch : patches ){
		std::string name { patch.name() };
		MPI_Bcast_H( name, part.comm(), part.comm().master() );
		if ( name != patch.name() ){
			std::cerr <<
			"Names of boundary patches must be equal on all mesh partitions!\n";
			MPI_Abort(part.comm(), MPI_ERR_OTHER);
		}
	}
	/* to find the limits for each patch among all partitions, each partition
	 * which has actual edges for that patch, needs to translate those edge
	 * indices to the global mesh, and then it's just a matter of finding the
	 * minimum and maximum index */
	ArrayXi patchLowerLimits ( nPatches );
	ArrayXi patchUpperLimits ( nPatches );
	for ( decltype( nPatches ) i{0}; i<nPatches; ++i ){
		if ( !(patches[i].empty()) ){
			/* translate first and last edge */
			patchLowerLimits(i) = edgeLocalToGlobal(
				patches[i].edges()[0],
				part.meshSizeGlobal(), part.meshSize(), part.offset()
			);
			patchUpperLimits(i) = edgeLocalToGlobal(
				patches[i].edges().tail(1)[0],
				part.meshSizeGlobal(), part.meshSize(), part.offset()
			);
		} else {
			patchLowerLimits(i) = std::numeric_limits<Index>::max();
			patchUpperLimits(i) = std::numeric_limits<Index>::min();
		}
	}
	MPI_Request req[2];
	MPI_Iallreduce(MPI_IN_PLACE, patchLowerLimits.data(), nPatches,
		MPI_INDEX, MPI_MIN, part.comm(), req );
	MPI_Iallreduce(MPI_IN_PLACE, patchUpperLimits.data(), nPatches,
		MPI_INDEX, MPI_MAX, part.comm(), req+1 );
	
	MPI_Waitall( 2, req, MPI_STATUSES_IGNORE );

	Array2Xi patchLimits (2, nPatches);
	patchLimits.row(0) = patchLowerLimits;
	patchLimits.row(1) = patchUpperLimits;
	return patchLimits;
}



template<typename MeshType>
void distributeBoundaryPatches(
	StructMeshPart<MeshType>& part,
	MeshType& mesh
){
	/* for convenient r/w access */
	std::vector<BoundaryPatch>& patches {
		mesh.boundaryPatches()
	};

	decltype( patches.size() ) nPatches;
	if ( part.comm().isMaster() ) nPatches = patches.size();

	MPI_Bcast_H( nPatches, part.comm(), part.comm().master() );

	Array2Xi patchLimits (2, nPatches);

	if ( !part.comm().isMaster() ) patches.resize(nPatches);
	
	for ( decltype(nPatches) i{0}; i<nPatches; ++i ){
		/* this could use Ibcast */
		MPI_Bcast_H( patches[i].name(), part.comm(), part.comm().master() );
		/* due to the restrictions on patches, only first and last edge index
		 * are needed */
		if ( part.comm().isMaster() )
			patchLimits.col(i) = Array2i{
				patches[i].edges()[0],
				patches[i].edges().tail(1)[0]
			};

		MPI_Bcast(
			patchLimits.data()+2*i, 2, MPI_INDEX,
			part.comm().master(), part.comm()
		);

		setLocalPatch( patches[i], part, mesh, patchLimits.col(i) );
	}

	part.boundaryPatchLimits( std::move(patchLimits) );
}



template<typename MeshType>
void redistributeBoundaryPatches(
	StructMeshPart<MeshType>& part,
	MeshType& mesh
){
	Index nPatches { part.boundaryPatchLimits().cols() };

	/* make sure the patch limits are available, otherwise make them available*/
	if (nPatches == 0)
		distributeBoundaryPatches(part, mesh);
	
	for ( Index i{0}; i<nPatches; ++i )
		setLocalPatch( mesh.boundaryPatches()[i],
			part, mesh, part.boundaryPatchLimits().col(i) );
}



} // end namespace hms
